'''

流程已合并到UnetOnnxNet.py中，此文件不再使用

'''

import argparse
from UnetOnnxNet import UnetOnnx
import time
import os
import cv2

model = UnetOnnx("v1_onnx.onnx")

def run_one(img):
    start_time = time.perf_counter()
    output = model.infer(img)
    end_time = time.perf_counter()
    print(f"infer file:{os.path.basename(img)},total time cost:{end_time - start_time}")
    return output

def run(input_pth: str, output_pth: str, mode:str):
    for file in os.listdir(input_pth):
        file_pth = os.path.join(input_pth, file)
        output = run_one(file_pth)
        new_name = file.replace(file.split('.')[-1], 'png')
        output_path = os.path.join(output_pth, new_name)
        cv2.imwrite(output_path, output)
        print(f"file:{file} saved in {output_path}")

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-i', type=str, required=True, help="input_dir")
    parser.add_argument('-o', type=str, required=True, help="output_dir")
    parser.add_argument('-m', type=str, required=True, help="fast or quality")
    args = parser.parse_args()
    if args.i == ' 'or args.o == ' ':
        print("input_dir or output_dir NOT FOUND")
    run(args.i, args.o, args.m)


